import csv
import time
import pandas as pd
import numpy as np
import os
import quantstats as qs
from itertools import groupby
# from tqdm import tqdm
from itertools import product
from multiprocessing import Pool
from functools import reduce
from pprint import pprint

pd.set_option('max_colwidth', 200)
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)
"""
股指期货日内短线突破
Created on 2024/05/23

@author: zhuky

# 交易时间段
 - 跨日收益率下降很多，每日第一段负收益居多


- 数据
    - 期货品种，价格（趋势突破判断，以及交易价格），期货合约乘数，保证金比率，期货交易手续费，期货履约手续费 (TBD: 不同价格的确认)
    - IF，分钟VWAP，300，每次开仓一手，手续费单边万分之0.4，保证金30%，止损1%，1min K线按照买一价卖一价成交 (TBD: 开仓数，不同品种)
- 评估指标
    - 1累计收益率，1年化收益率
    - 0交易次数，0盈利次数，0亏损次数，0胜率
    - [SIGN gain-loss]0平均收益-盈亏加一起，0平均单笔盈利-只看盈利段，0平均单笔亏损-只看亏损段，0盈亏比，0最大单笔盈利，0最大单笔亏损，0最大连续盈利次数，0最大连续亏损次数，0平均连续盈利次数，0平均连续亏损次数
    - 0总的毛盈利，0总的毛亏损，1最大回撤，1最长再创新高周期，1夏普比率
    - 开多仓/开空仓 0平均每笔收益
    - 0单笔收益率分布
    - 1单月收益率时序图
- 参数：日内时间区间划分P，波幅大小控制W (TBD: P,W时序可变)
- 参数优化指标：累计收益率/最大回撤
- 策略过程
    - 获取每日品种的交易价格序列
        - 趋势跟踪所使用的价格：price(TBD: 时段均价) - sensitive, vwap 
        - 趋势交易所使用的价格：price(TBD: 对手方价格)
        - 趋势上下轨所使用的价格：price(TBD: 上一时段最高价(最高价的最高价)，上一时段最低价(最低价的最低价))；
    - 拆分交易时段，每个时段存储：这个时段的价格序列，上个时段的最高价、最低价  (TBD: 存在交易时段不连续情况，考虑删除) - TIME
    - 每个交易时段最多交易一手持有到周期结束，存储每个时段的交易次数、成本、盈亏额
    - 样本内网格调优P,W，样本外回测结果 (TBD: 滚动调优)
        - 全数据集：20110104 - 20240521
        - 样本内：20110104 - 20210104
        - 样本外：20210105 - 20240521

    - ! 参数稳健性
    - 结合 dual ,  open - 新算法
"""
"""
TODO:
 - [1]价格敏感性测试，price -> vwap
 - [0]滚动调优, 10 -> 10
"""

test_folder = f'test_basic/{time.strftime("%m%d", time.localtime())}/highFeelowMargin/'
os.makedirs(test_folder, exist_ok=True)
future = 'IF'
multiplier = 300
fee_ratio = 2.53e-4
margin_ratio = 0.15

max_loss_rate = -0.01

process_num = 16

global start_index, end_index
global start_date, end_date

with pd.HDFStore(f"dataset/{future}.h5", 'r') as store:
    dates = list(store.keys())
    len_dates = len(dates)
    # test_len = int(len_dates / 3)
    # start_index = int(len_dates /3) * 2
    start_index = 0
    end_index = len_dates  # 不包含end_index
    # start_index = 0
    # end_index = dates.index('/20140303')  # 不包含end_index

    start_date = dates[start_index][1:]
    end_date = dates[end_index - 1][1:]

note = 'Basic_hFlM_sharpe_all'

global op_file
op_file = f'output_{start_date}_{end_date}_{note}.csv'
# print(op_file)


def count_consecutive_tf(data):
    """
    用于计数连续盈利or亏损的次数
    :param data: 元素为True or False的数组
    :return:
    """
    # 使用groupby来分组连续的True或False
    grouped = groupby(data, key=lambda x: x)
    consecutive_true_counts = []
    for is_true, group in grouped:
        if is_true:  # 如果当前分组是True
            count = sum(1 for _ in group)  # 计算True的数量
            consecutive_true_counts.append(count)
    return consecutive_true_counts

    # TBD: 可能得在观察到第一个突破的下一分钟交易，目前是观察到第一个突破，以该突破价成交; 是否中途需要补充保证金
    # 止损线，开空反而价格上涨有问题
    # 交易价格：信号发生的下一个价格


def _cal_ret_ratio(price_list, index, train, date, num, direction=1):
    """
    计算收益率，已知突破点位和开仓方向时
    :param price_list: 价格序列
    :param index: 突破后的交易价格索引
    :param train: True则只返回收益率，False返回所有指标
    :param direction: long=1,short=-1
    :return:
    """
    val = price_list[index]
    end_num = next((x for x in price_list[index:] if x < val * (1 + max_loss_rate * margin_ratio)),
                   price_list[-1]) if direction + 1 else next(
        (x for x in price_list[index:] if x > val * (1 - max_loss_rate * margin_ratio)),
        price_list[-1])
    gain = direction * (end_num - val) * multiplier - val * multiplier * fee_ratio
    ret_ratio = gain / (val * multiplier * margin_ratio)
    if train:
        return pd.Series([date, num, ret_ratio], index=['date', 'num', 'ret_ratio'])
    res = pd.Series([date, num, index, val, end_num, direction, gain, ret_ratio],
                    index=['date', 'num', 'rush_index', 'rush_num', 'tail_num', 'direction', 'gain', 'ret_ratio'])
    return res


def _gain_loss_series(series, train):
    """
    针对趋势价格，上轨、下轨，给出当前时间段的交易机会和收益情况
    :param series:
    :param train: True则只返回收益率，False返回所有指标
    :return: 突破时点，交易价格，平仓价格，多空方向，收益，收益率
    """
    g_index = next(
        (index for index, value in enumerate(series['price_list']) if value >= series['upper_lower_bound'][0]), np.inf)
    g_index = g_index + 1 if g_index is not np.inf and g_index < len(series['price_list']) - 1 else np.inf
    l_index = next(
        (index for index, value in enumerate(series['price_list']) if value <= series['upper_lower_bound'][1]), np.inf)
    l_index = l_index + 1 if l_index is not np.inf and l_index < len(series['price_list']) - 1 else np.inf
    min_g_l = min(g_index, l_index)
    if min_g_l is np.inf:
        return pd.Series()
    
    direction = 1 if min_g_l == g_index else -1
    return _cal_ret_ratio(series['price_list'], min_g_l, train, series['date'], series['num'], direction)


def valuate_param(daily_ret):
    """
    指标：calmar - 基于累加收益率
    """
    daily_ret.index = pd.to_datetime(daily_ret.index)
    cumr = qs.stats.cumr(daily_ret)
    max_dd = qs.stats.max_drawdown(daily_ret,compounded=False)
    calmar = qs.stats.calmar(daily_ret,compounded=False)
    return cumr, max_dd, calmar


def get_param_valuation_backtest(W, P, train=True):
    """
    数据切片，整合出收益情况
    :param data: 分钟K线数据 on date
    :param W: 策略参数区间宽度W
    :param P: 策略参数区间长度P
    :param train: True只计算收益率，False计算所有指标
    :return:
    """

    def get_daily_price_split(date, daily_data, P):
        # NOTE: 不可以使用resample,其按照时间补全重采样
        data_use_daily = pd.DataFrame(columns=['date', 'num', 'price_list'])
        num = 0
        for i in range(0, daily_data.shape[0], P):
            num += 1
            j = i + P if i + P < daily_data.shape[0] else daily_data.shape[0]
            data_period = daily_data.iloc[i:j]
            price_list = data_period['price'].values.tolist()
            period_dict = {'date': date[1:], 'num': num, 'price_list': price_list}
            data_use_daily = data_use_daily._append(period_dict, ignore_index=True)
        return data_use_daily

    with pd.HDFStore(f"dataset/{future}.h5", 'r') as store:
        dates = list(store.keys())
        df_list = [get_daily_price_split(date, store.get(date), P) for date in dates[start_index:end_index]]

    data_use = pd.concat(df_list, ignore_index=True)
    data_use.sort_values(by=['date', 'num'], ascending=True, inplace=True)
    # print(data_use.shape)
    data_use.loc[:, 'max_min_cur'] = data_use['price_list'].apply(lambda x: (max(x), min(x)))
    # test: 不考虑跨日的结果
    data_use.loc[:, 'max_min_pre'] = data_use.groupby('date')['max_min_cur'].shift(1)
    # 考虑跨日的结果(亏损增加)
    # data_use.loc[:, 'max_min_pre'] = data_use.loc[:, 'max_min_cur'].shift(1)
    data_use.dropna(inplace=True)
    # print(data_use.shape)

    data_use.loc[:, 'upper_lower_bound'] = data_use['max_min_pre'].apply(
        lambda x: (x[0] + W * (x[0] - x[1]), x[1] - W * (x[0] - x[1])))

    res_df = data_use.apply(lambda x: _gain_loss_series(x, train), axis=1).dropna(how='all')
    # print(res_df.shape)

    # test_sample = pd.merge(data_use,res_df, 'left', on=['date','num'])
    # test_sample.to_csv(test_folder + f'testSample_{W}_{P}_{start_date}_{end_date}_{note}.csv',index=None)

    res_daily = res_df.groupby('date', as_index=True)['ret_ratio'].sum()
    res_daily.index = pd.to_datetime(res_daily.index)
    # print(res_daily.shape)

    res_daily.to_csv(test_folder + f'dailyRet_{W}_{P}_{start_date}_{end_date}_{note}.csv')

    if train:

        cumr, max_dd, calmar = valuate_param(res_daily)
        print(val_dict:={
            'W': W,
            'P': P,
            'cumr': cumr,
            'max_drawdown': max_dd,
            'calmar': calmar
        })        
    else:
        # 累计收益率
        cumr = qs.stats.cumr(res_daily)
        # 年化收益率
        cagr = qs.stats.cagr(res_daily,compounded=False)
        # 最大回撤
        max_drawdown = qs.stats.max_drawdown(res_daily,compounded=False)
        # 累计收益率/最大回撤
        cumr_md = cumr / abs(max_drawdown)
        # calmar
        calmar = qs.stats.calmar(res_daily,compounded=False)
        # 夏普比率
        sharpe = qs.stats.sharpe(res_daily)
        # 交易次数
        trade_num = res_df['direction'].abs().sum()
        # 盈利次数
        gain_num = (res_df['gain'] > 0).sum()
        # 亏损次数
        loss_num = (res_df['gain'] < 0).sum()
        # 胜率
        win_rate = gain_num / trade_num
        # 总的毛盈利，总的毛亏损
        gain_amount = res_df.loc[res_df['gain'] > 0, 'ret_ratio'].sum()
        loss_amount = res_df.loc[res_df['gain'] < 0, 'ret_ratio'].sum()
        # 平均收益，平均单笔盈利，平均单笔亏损
        avg_ret_amount_per_trade = res_daily.sum() / trade_num
        avg_gain_amount_per_gain = gain_amount / gain_num
        avg_loss_amount_per_loss = loss_amount / loss_num
        # 最大单笔盈利，最大单笔亏损
        max_gain_amount = res_df['ret_ratio'].max()
        max_loss_amount = res_df['ret_ratio'].min()
        # 最大连续盈利次数，最大连续亏损次数
        consecutive_gain_counts = count_consecutive_tf(res_df['gain'] > 0)
        consecutive_loss_counts = count_consecutive_tf(res_df['gain'] < 0)
        max_consecutive_gain = max(consecutive_gain_counts) if len(consecutive_gain_counts) else None
        max_consecutive_loss = max(consecutive_loss_counts) if len(consecutive_loss_counts) else None
        # 平均连续盈利次数，平均连续亏损次数
        avg_consecutive_gain = np.mean(consecutive_gain_counts)
        avg_consecutive_loss = np.mean(consecutive_loss_counts)
        # 开多每笔收益，开空每笔收益
        avg_ret_amount_per_long = np.mean(res_df.loc[res_df['direction'] == 1, 'ret_ratio'].values)
        avg_ret_amount_per_short = np.mean(res_df.loc[res_df['direction'] == -1, 'ret_ratio'].values)
        # EV
        ev = win_rate * avg_gain_amount_per_gain + (1 - win_rate) * avg_loss_amount_per_loss
        # # 单月收益率
        # monthly_returns = qs.stats.monthly_returns(test_sample['日收益率'])

        # qs.plots.snapshot(test_sample['日收益率'], title=f'IF P={P} W={W}', show=True)
        qs.reports.html(res_daily, compounded=False, strategy_title=f'{future}_IntraBreak', title=f'{future} P={P} W={W}',
                        output=test_folder + f'{future}_{W}_{P}_{start_date}_{end_date}_{note}.html')
        val_dict = {
            'W': W,
            'P': P,
            'calmar':calmar,
            'cumr-md': cumr_md,
            '累计收益率': cumr,
            '年化收益率': cagr,
            '最大回撤': max_drawdown,
            '夏普比率': sharpe,
            '交易次数': trade_num,
            '盈利次数': gain_num,
            '亏损次数': loss_num,
            '胜率': win_rate,
            '总的毛盈利': gain_amount,
            '总的毛亏损': loss_amount,
            '平均收益': avg_ret_amount_per_trade,
            '平均单笔盈利': avg_gain_amount_per_gain,
            '平均单笔亏损': avg_loss_amount_per_loss,
            '最大单笔盈利': max_gain_amount,
            '最大单笔亏损': max_loss_amount,
            '最大连续盈利次数': max_consecutive_gain,
            '最大连续亏损次数': max_consecutive_loss,
            '平均连续盈利次数': avg_consecutive_gain,
            '平均连续亏损次数': avg_consecutive_loss,
            '开多每笔收益': avg_ret_amount_per_long,
            '开空每笔收益': avg_ret_amount_per_short,
            # '单月收益率':monthly_returns,
            'EV': ev,
        }
        print({
            'W': W,
            'P': P,
            'calmar': calmar,
            'EV': ev
        })
        # pprint(val_dict)

    with open(test_folder + op_file, 'a', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(list(val_dict.values()))

    return val_dict

def run_single_valuation(param_tuple):
    W,P = param_tuple
    get_param_valuation_backtest(W,P,True)


def run_full_valuation(param_tuple):
    W,P = param_tuple
    get_param_valuation_backtest(W,P,False)


# 样本内：已有每日收益率 \to 样本内累计收益率/最大回撤
# 参数优化

if __name__ == '__main__':
    start = time.time()
    print('run start:', time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
    
    W_list = np.arange(0.2, 0.9, 0.02)
    P_list = np.arange(20, 101, 5)

    W = 0.48
    P = 85

    value_all = True

    param_res_col = ['W', 'P', 'cumr', 'max_drawdown', 'calmar'] if not value_all else ['W', 'P', 'calmar',
                                                                                        'cumr_md',
                                                                                                  '累计收益率',
                                                                                                  '年化收益率',
                                                                                                  '最大回撤',
                                                                                                  '夏普比率',
                                                                                                  '交易次数',
                                                                                                  '盈利次数',
                                                                                                  '亏损次数', '胜率',
                                                                                                  '总的毛盈利',
                                                                                                  '总的毛亏损',
                                                                                                  '平均收益',
                                                                                                  '平均单笔盈利',
                                                                                                  '平均单笔亏损',
                                                                                                  '最大单笔盈利',
                                                                                                  '最大单笔亏损',
                                                                                                  '最大连续盈利次数',
                                                                                                  '最大连续亏损次数',
                                                                                                  '平均连续盈利次数',
                                                                                                  '平均连续亏损次数',
                                                                                                  '开多每笔收益',
                                                                                                  '开空每笔收益', 'EV']
    param_res_df = pd.DataFrame(columns=param_res_col)
    with open(test_folder + op_file, 'w', newline='') as file:
        writer = csv.writer(file)
        writer.writerow(param_res_col)

    run_full_valuation((W,P))

    # combinations = list(product(W_list, P_list))
    # with Pool(processes=process_num) as pool:
    #     pool.map(run_single_valuation, combinations) if not value_all else pool.map(run_full_valuation, combinations)

    print(f'running: {round(time.time()-start,2)}s')
